# 题目1
# 假设有一组视频数据，每个数据有2个特征，分别是播放量，点赞数，对应特征的标签label为收入
# 制作一个模型，通过输入播放量和点赞数，预测视频的收入
# 具体要求如下：
# 1. 随机20组数据
# 2. 播放量应该在90~100之间的整数
# 3. 点赞数应该为当前视频播放量的5分之一再加上正负百分之5的误差
# 4. 收入公式为(期望函数）：y=0.1*播放量+0.3*点赞数
# 5. 将数据用散点的方式排列到三维空间中，用画图的方式展示训练过程
#
# 注意：这个练习可能会出现问题，原因在于没有归一化，你可以手动归一化


import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 生成数据
np.random.seed(42)
torch.manual_seed(42)

# 随机生成20组数据
n_samples = 20

# 播放量在90~100之间的整数
play_count = torch.randint(90, 101, (n_samples, 1)).float()

# 点赞数为播放量的1/5加上正负5%的误差
base_likes = play_count / 5
error_range = base_likes * 0.05
likes = base_likes + torch.randn(n_samples, 1) * error_range

# 收入公式：y = 0.1*播放量 + 0.3*点赞数
income = 0.1 * play_count + 0.3 * likes

# 合并特征
X = torch.cat([play_count, likes], dim=1)
y = income

print("数据示例：")
print("播放量\t点赞数\t收入")
for i in range(5):
    print(f"{X[i, 0].item():.1f}\t{X[i, 1].item():.1f}\t{y[i].item():.2f}")

# 2. 数据归一化
X_mean = X.mean(dim=0)
X_std = X.std(dim=0)
y_mean = y.mean()
y_std = y.std()

X_normalized = (X - X_mean) / X_std
y_normalized = (y - y_mean) / y_std


# 3. 创建模型
class VideoIncomeModel(nn.Module):
    def __init__(self):
        super(VideoIncomeModel, self).__init__()
        self.linear = nn.Linear(2, 1)

    def forward(self, x):
        return self.linear(x)


model = VideoIncomeModel()

# 4. 设置超参数
learning_rate = 0.01
epochs = 1000
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# 5. 训练过程记录
losses = []
predictions_history = []

# 6. 训练模型
for epoch in range(epochs):
    # 前向传播
    y_pred = model(X_normalized)

    # 计算损失
    loss = loss_fn(y_pred, y_normalized)
    losses.append(loss.item())

    # 反向传播
    loss.backward()

    # 更新参数
    optimizer.step()

    # 清零梯度
    optimizer.zero_grad()

    # 每100轮记录一次预测结果用于可视化
    if epoch % 100 == 0:
        with torch.no_grad():
            pred_denormalized = y_pred * y_std + y_mean
            predictions_history.append(pred_denormalized.numpy().flatten())

# 7. 最终预测结果
with torch.no_grad():
    final_pred = model(X_normalized)
    final_pred_denormalized = final_pred * y_std + y_mean

# 8. 可视化结果
fig = plt.figure(figsize=(15, 5))

# 8.1 显示训练损失
ax1 = fig.add_subplot(131)
ax1.plot(losses)
ax1.set_title('训练损失变化')
ax1.set_xlabel('训练轮次')
ax1.set_ylabel('损失值')

# 8.2 3D散点图显示数据和拟合平面
ax2 = fig.add_subplot(132, projection='3d')

# 原始数据点
ax2.scatter(X[:, 0].numpy(), X[:, 1].numpy(), y.numpy(), c='red', marker='o', label='真实数据')

# 模型预测平面
play_grid = np.linspace(X[:, 0].min(), X[:, 0].max(), 10)
likes_grid = np.linspace(X[:, 1].min(), X[:, 1].max(), 10)
Play, Likes = np.meshgrid(play_grid, likes_grid)

# 将网格数据归一化后预测
Play_norm = (Play - X_mean[0].item()) / X_std[0].item()
Likes_norm = (Likes - X_mean[1].item()) / X_std[1].item()
Income_pred = 0.1 * Play + 0.3 * Likes  # 理论公式

ax2.plot_surface(Play, Likes, Income_pred, alpha=0.3, color='blue')
ax2.set_xlabel('播放量')
ax2.set_ylabel('点赞数')
ax2.set_zlabel('收入')
ax2.set_title('视频数据3D可视化')

# 8.3 训练过程动画效果
ax3 = fig.add_subplot(133)
ax3.scatter(X[:, 0].numpy(), y.numpy(), c='red', label='真实数据')

# 绘制训练过程中不同阶段的预测结果
colors = ['lightblue', 'lightgreen', 'lightyellow', 'lightpink', 'lightgray']
for i, pred in enumerate(predictions_history):
    if i < len(colors):
        ax3.scatter(X[:, 0].numpy(), pred, c=colors[i], alpha=0.7,
                    label=f'第{i * 100}轮预测')

# 最终预测结果
ax3.scatter(X[:, 0].numpy(), final_pred_denormalized.numpy(), c='blue',
            label='最终预测', marker='x')

ax3.set_xlabel('播放量')
ax3.set_ylabel('收入')
ax3.set_title('训练过程可视化')
ax3.legend()

plt.tight_layout()
plt.show()

# 9. 模型评估
with torch.no_grad():
    mse = torch.mean((final_pred_denormalized - y) ** 2)
    print(f"\n模型评估:")
    print(f"最终损失: {losses[-1]:.6f}")
    print(f"均方误差: {mse.item():.6f}")

    # 显示学习到的参数
    weights = model.linear.weight.data
    bias = model.linear.bias.data
    print(
        f"学习到的权重: 播放量权重={weights[0, 0].item() * y_std / X_std[0].item():.4f}, 点赞数权重={weights[0, 1].item() * y_std / X_std[1].item():.4f}")
    print(f"理论权重: 播放量权重=0.1, 点赞数权重=0.3")

# 10. 测试新数据
print("\n测试新数据:")
test_play = torch.tensor([[95.0, 19.0]])  # 播放量95，点赞数19
test_play_normalized = (test_play - X_mean) / X_std
with torch.no_grad():
    test_pred_normalized = model(test_play_normalized)
    test_pred = test_pred_normalized * y_std + y_mean
    theoretical_income = 0.1 * 95 + 0.3 * 19  # 理论收入
    print(f"播放量: 95, 点赞数: 19")
    print(f"模型预测收入: {test_pred.item():.2f}")
    print(f"理论收入: {theoretical_income:.2f}")
